package cn.com.ry.framework.monitormeter.metric.jdbc;

import cn.com.ry.framework.monitormeter.instrument.linklog.LinkLogMeterRegistry;
import cn.com.ry.framework.monitormeter.util.Assert;
import io.micrometer.core.instrument.Gauge;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import io.micrometer.core.instrument.binder.BaseUnits;
import io.micrometer.core.instrument.binder.MeterBinder;
import io.micrometer.prometheus.PrometheusMeterRegistry;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.function.ToDoubleFunction;

public class DBTableMetrics implements MeterBinder {
    private final DataSource dataSource;
    private final String dataSourceName;
    private final Iterable<Tag> tags;
    private List<JDBCBean> jdbcBeanList = new ArrayList<JDBCBean>();


    public DBTableMetrics(DataSource dataSource, String dataSourceName, List<JDBCBean> jdbcBeanList, Iterable<Tag> tags) {
        this.dataSourceName = dataSourceName;
        this.dataSource = dataSource;
        this.jdbcBeanList = jdbcBeanList;
        this.tags = tags;
    }

    public DBTableMetrics(DataSource dataSource, String dataSourceName, List<JDBCBean> jdbcBeanList) {
        this(dataSource, dataSourceName, jdbcBeanList, Tags.of("remark", "数据库监控"));
    }

    @Override
    public void bindTo(MeterRegistry registry) {
        if (registry != null) {
            Assert.notNull(jdbcBeanList, "监控的sql不能为空");
            if (jdbcBeanList != null) {
                if(registry instanceof LinkLogMeterRegistry){
                    for (JDBCBean jdbcBean : jdbcBeanList) {
                        String sql = jdbcBean.getSql();
                        String name = jdbcBean.getName();
                        ToDoubleFunction<DataSource> totalRows = ds -> {
                            try (Connection conn = ds.getConnection();
                                 PreparedStatement ps = conn.prepareStatement(sql);
                                 ResultSet rs = ps.executeQuery()) {
                                rs.next();
                                return rs.getInt(1);
                            } catch (SQLException ignored) {
                                return 0;
                            }
                        };
                        Gauge.builder("db.query." + dataSourceName + "." + name + ".size", dataSource, totalRows)
                                .tags(tags)
                                .description("Number of rows in a database table")
                                .baseUnit(BaseUnits.ROWS)
                                .register(registry);
                    }
                }else if(registry instanceof PrometheusMeterRegistry){
                    for (JDBCBean jdbcBean : jdbcBeanList) {
                        String sql = jdbcBean.getSql();
                        String name = jdbcBean.getName();
                        ToDoubleFunction<DataSource> totalRows = ds -> {
                            try (Connection conn = ds.getConnection();
                                 PreparedStatement ps = conn.prepareStatement(sql);
                                 ResultSet rs = ps.executeQuery()) {
                                rs.next();
                                return rs.getInt(1);
                            } catch (SQLException ignored) {
                                return 0;
                            }
                        };
                        Gauge.builder("db.query.size", dataSource, totalRows)
                                .tags(Tags.of(Tag.of("datasource.name",dataSourceName),Tag.of("table.name",name)))
                                .description("Number of rows in a database table")
                                .baseUnit(BaseUnits.ROWS)
                                .register(registry);
                    }
                }


            }
        }
    }
}
